import numpy as np
import pandas as pd

import os
import pickle
from sklearn.metrics import f1_score, roc_auc_score, roc_curve, auc, accuracy_score

# results_folder_path = "/media/hsmy/N033/research/"
results_folder_path = "/media/hsmy/wanghao_18T/tj/results_202510/"
task_excel_data = [["task", "model", "acc", "auc",  "ci", "ci-lower", "ci-upper", "f1", "f1s", "tpr", "fpr",]]
for root, dirs, files in os.walk(results_folder_path):
    for _file in files:
        if _file.endswith(".pkl"):
            exp_name = os.path.basename(root)
            exp_name = exp_name.replace("ours_", "ours")
            exp_name_arr = exp_name.split("_")
            _task = exp_name_arr[1]
            _model = exp_name_arr[6]
            # _task = exp_name_arr[1] + exp_name_arr[2] + exp_name_arr[3]
            # _model = exp_name_arr[8]
            _pkl_path = os.path.join(root, _file)
            _gt_arr = []
            _pred_arr = []
            _pred_prob_arr = []
            _pred_prob_arr2 = []
            with open(_pkl_path, 'rb') as f:
                _pkl_data = pickle.load(f)
                for _key in _pkl_data.keys():
                    _item = _pkl_data[_key]
                    _gt = _item['label']
                    _prob = _item['prob'][0]
                    _pred_arr.append(np.argmax(_prob))
                    _gt_arr.append(_gt)
                    _pred_prob_arr.append(_prob[1])
                    _pred_prob_arr2.append(_prob)

            _accuracy = accuracy_score(_gt_arr, _pred_arr)

            is_multi_class = np.unique(_gt_arr).size > 2
            if is_multi_class:
                fpr, tpr = 0, 0  # ROC曲线本身在多分类中不直接可视化，可设为None或后续处理
                _auc = roc_auc_score(_gt_arr, _pred_prob_arr2, multi_class='ovr', average='macro')
            else:
                fpr, tpr, _ = roc_curve(_gt_arr, _pred_prob_arr)
                _auc = auc(fpr, tpr)

            indices = np.arange(len(_gt_arr))
            rng = np.random.RandomState(42)
            aurocs = []
            f1s = []
            _gt_arr = np.array(_gt_arr)
            _pred_arr = np.array(_pred_arr)
            _pred_prob_arr = np.array(_pred_prob_arr)
            _pred_prob_arr2 = np.array(_pred_prob_arr2)
            for _ in range(1000):
                sample_indices = rng.choice(indices, size=len(_gt_arr), replace=True)
                gt_sample = _gt_arr[sample_indices]
                pred_sample = _pred_arr[sample_indices]
                if len(np.unique(gt_sample)) < 2:
                    continue
                if is_multi_class:
                    f1 = f1_score(gt_sample, pred_sample, average='weighted')
                else:
                    f1 = f1_score(gt_sample, pred_sample)

                f1s.append(f1)
            f1s = np.array(f1s)
            lower = np.percentile(f1s, 2.5)
            upper = np.percentile(f1s, 97.5)
            mean_f1 = np.mean(f1s)
            ci = (upper - lower) / 2
            task_excel_data.append(
                [_task, _model, f"{_accuracy:.3f}", f"{_auc:.3f}", f"{ci:.3f}", f"{lower:.3f}", f"{upper:.3f}", f"{mean_f1:.3f}",
                str(f1s.tolist()), str(np.array(tpr).tolist()), str(np.array(fpr).tolist())])
            print(f"{_task} {_model} done")

pd.DataFrame(task_excel_data).to_excel("results.xlsx", index=None, header=None)
pd.DataFrame(task_excel_data).to_csv("results.csv", index=None, header=None)
